"""
参考
https://jacco.ompf2.com/2019/12/11/probability-theory-for-physically-based-rendering/
"""
import numpy as np
import taichi as ti

vec2 = ti.math.vec2
vec3 = ti.math.vec3
EPS = 1e-3
INF = 1e10

@ti.dataclass
class Ray:
    """
    射线,起点original,方向direct
    """
    original : vec3
    direct   : vec3

@ti.dataclass
class Sphere:
    """
    球中心center,半径radius
    material 材质索引
    """
    center: vec3
    radius: ti.f32
    x_axis : vec3
    y_axis : vec3
    z_axis : vec3

@ti.dataclass
class Trangle:
    """
    三角形顶点a,b,c,法线方向n
    material 材质索引
    """
    a : vec3
    b : vec3
    c : vec3
    a_uv : vec2
    b_uv : vec2
    c_uv : vec2
    n : vec3

@ti.dataclass
class Material:
    """
    为了方便光被认为是一种材质
    材质color颜色,alpha表面粗糙度0~1
    """
    type  : ti.u32   #MaterialType
    color : vec3
    tex  : ti.i32   #材质索引-1表示不存在
    a   : ti.f32    #表面粗糙度roughness,如果是光表示强度
    f   : vec3      #菲涅尔参数
    metallic : ti.f32 #金属度
    mu   : ti.f32   #材质折射率

@ti.dataclass
class Light:
    type    :ti.u32 #LightType
    power   :ti.f32 #灯的功率
    tex     :ti.i32
    color   :vec3   #灯的颜色
    p       :vec3   #灯的位置
    dir     :vec3   #方向
    angle   :ti.f32 #聚光灯角度
    offset  :ti.i32 #发光体三角形偏移
    count   :ti.i32
    cdf_v : ti.f32
    pdf_v : ti.f32

@ti.dataclass
class SphereBound:
    """
    包围球
    """
    type : ti.u32 #PrimitiveType
    material : ti.i32
    center : vec3
    radius : ti.f32
    x_axis : vec3
    y_axis : vec3
    z_axis : vec3
    offset : ti.i32 #对应的三角形列表的偏移位置
    count  : ti.i32 #三角形数量

@ti.dataclass
class AabbBound:
    """
    轴对齐包围盒
    """
    type   : ti.u32  #PrimitiveType
    min     :vec3
    max     :vec3  #aabb包围盒
    offset : ti.i32 #对应的三角形列表的偏移位置或者表示使用三角形
    count  : ti.i32 #三角形数量
    material:ti.i32

@ti.dataclass
class Intersection:
    """
    交点信息
    """
    type:ti.u32 #PrimitiveType
    t   :ti.f32 #距离
    p   :vec3   #交点位置
    n   :vec3   #法线
    m   :ti.i32 #材质索引
    i   :ti.i32 #三角形索引或者球的索引

@ti.func
def radical_inverse_vdc(bits:ti.u32):
    """
    Van Der Corput序列
    """
    bits = (bits << 16) | (bits >> 16)
    bits = ((bits & ti.u32(0x55555555)) << 1) | ((bits & ti.u32(0xAAAAAAAA)) >> 1)
    bits = ((bits & ti.u32(0x33333333)) << 2) | ((bits & ti.u32(0xCCCCCCCC)) >> 2)
    bits = ((bits & ti.u32(0x0F0F0F0F)) << 4) | ((bits & ti.u32(0xF0F0F0F0)) >> 4)
    bits = ((bits & ti.u32(0x00FF00FF)) << 8) | ((bits & ti.u32(0xFF00FF00)) >> 8)
    return float(bits) * 2.3283064365386963e-10

@ti.func
def hammersley(i:ti.u32,n:ti.u32)->vec2:
    """
    低方差2d随机序列
    a low-discrepancy sequence
    https://learnopengl.com/PBR/IBL/Specular-IBL
    """
    return ti.Vector([float(i)/float(n),radical_inverse_vdc(i)])

@ti.func
def random_vec2()->vec2:
    return ti.Vector([ti.random(),ti.random()])

@ti.func
def importance_sample_ggx(xi:vec2,n:vec3,roughness:float)->vec3:
    """
    ggx采样
    """
    a = roughness * roughness
    phi = 2. * np.pi * xi.x
    cos_theta = ti.math.sqrt((1.-xi.y)/(1.+(a * a-1.) * xi.y))
    sin_theta = ti.math.sqrt(1.-cos_theta)
    h = ti.Vector([ti.math.cos(phi) * sin_theta,ti.math.sin(phi) * sin_theta,cos_theta])
    up = ti.Vector([0.,0.,1.]) if ti.abs(n.z) < .999 else ti.Vector([1.,0.,0.])
    tangent = ti.math.normalize(ti.math.cross(up,n))
    bitangent = ti.math.cross(n,tangent)
    sample_vec = tangent * h.x + bitangent * h.y + n * h.z
    return ti.math.normalize(sample_vec)

@ti.func
def random_dir(n):
    """
    平面的法线为n,在平面一侧的随机方向
    """
    u = ti.Vector([1.0, 0.0, 0.0])
    if ti.abs(n[1]) < 1 - 1e-3:
        u = n.cross(ti.Vector([0.0, 1.0, 0.0])).normalized()
    v = n.cross(u)
    phi = 2 * np.pi * ti.random(ti.f32)
    r = ti.random(ti.f32)
    ay = ti.sqrt(r)
    ax = ti.sqrt(1 - r)
    return ax * (ti.cos(phi) * u + ti.sin(phi) * v) + ay * n

@ti.func
def sphere_random_dir(n):
    """
    随机方向
    """
    u = ti.Vector([1.0, 0.0, 0.0])
    if ti.abs(n[1]) < 1 - 1e-3:
        u = n.cross(ti.Vector([0.0, 1.0, 0.0])).normalized()
    v = n.cross(u)
    alpha = 2 * np.pi * ti.random(ti.f32)
    theta = .5 * np.pi * ti.random(ti.f32)

    return (u * ti.cos(alpha) + v * ti.sin(alpha)) * ti.cos(theta) + n * ti.sin(theta)

@ti.func
def circle_random_samples(r):
    """
    圆形采样
    """
    r_0 = r * ti.math.sqrt(ti.random())
    phi = 2. * np.pi * ti.random()
    return ti.Vector([r_0 * ti.math.sin(phi),r_0 * ti.math.cos(phi)])

@ti.func
def cross_ray_sphere(ray:Ray,sphere:ti.template()):
    """
    射线和球求交点,
    返回参数t,交点p,和交点处的法线n
    """
    B = ray.original-sphere.center
    a = ti.math.dot(ray.direct,ray.direct)
    b = 2*ti.math.dot(ray.direct,B)
    c = ti.math.dot(B,B)-sphere.radius**2
    delta = b**2 - 4*a*c
    p = ti.Vector([0.,0.,0.])
    n = ti.Vector([0.,0.,1.])
    t = 0.
    isc = 0
    if delta>=0:
        delta_sq = ti.math.sqrt(delta)
        t1 = (-b + delta_sq)/2 * a
        t2 = (-b - delta_sq)/2 * a
        if t1>0 and t2>0:
            t = ti.math.min(t1,t2)
        elif t2>0:
            t = t2
        else:
            t = t1
        p = ray.direct * t + ray.original
        n = ti.math.normalize(p-sphere.center)
        isc = t>0
    return isc,t,p,n

@ti.func
def cross_ray_trangle(ray:Ray,trangle:Trangle):
    """
    射线和三角求交点,
    返回交点p和交点处的法线n
    """
    p = ti.Vector([0.,0.,0.])
    n = ti.Vector([0.,0.,1.])
    t = 0.
    isc = 0
    t = ti.math.dot(trangle.n,trangle.a-ray.original)/ti.math.dot(trangle.n,ray.direct)
    if t>=0:
        e1 = trangle.b - trangle.a
        e2 = trangle.c - trangle.a
        s = ray.original - trangle.a
        s1 = ti.math.cross(ray.direct,e2)
        s2 = ti.math.cross(s,e1)
        coeff = 1./ti.math.dot(s1,e1)
        b0 = coeff * ti.math.dot(s2,e2)
        if b0>=0:
            b1 = coeff * ti.math.dot(s1,s)
            if b1>=0:
                b2 = coeff * ti.math.dot(s2,ray.direct)
                if b2>=0 and (1.-b1-b2)>=0:
                    n = trangle.n
                    p = ray.original + ray.direct*t
                    isc = 1

    return isc,t,p,n

@ti.func
def ray_aabb_intersection(ray:Ray,box_min, box_max):
    """
    射线和包围盒是否相交
    """
    intersect = 1

    near_int = -INF
    far_int = INF

    for i in ti.static(range(3)):
        if ray.direct[i] == 0:
            if ray.original[i] < box_min[i] or ray.original[i] > box_max[i]:
                intersect = 0
        else:
            i1 = (box_min[i] - ray.original[i]) / ray.direct[i]
            i2 = (box_max[i] - ray.original[i]) / ray.direct[i]

            new_far_int = ti.max(i1, i2)
            new_near_int = ti.min(i1, i2)

            far_int = ti.min(new_far_int, far_int)
            near_int = ti.max(new_near_int, near_int)

    if near_int > far_int:
        intersect = 0
    return intersect#, near_int, far_int

@ti.func
def distribution_ggx(nh_dot,roughness):
    a = roughness * roughness
    a2 = a * a
    nh2_dot = nh_dot * nh_dot
    denom = nh2_dot * (a2 - 1.) + 1.
    return a2 / (np.pi * denom * denom)

@ti.func
def geometry_schlick_ggx(nv_dot,roughness):
    r = roughness - 1.
    k = (r * r)/8.
    return nv_dot/(nv_dot * (1.-k) + k)

@ti.func
def geometry_smith(nv_dot,nl_dot,roughness):
    return geometry_schlick_ggx(nv_dot,roughness) * geometry_schlick_ggx(nl_dot,roughness)

@ti.func
def fresnel_schlick(nv_dot,f0):
    """
    菲涅尔近似公式
    """
    return f0 + (1. - f0) * ti.math.pow(1.-nv_dot,5.)

@ti.func
def importance_sample_ggx_pdf(wi,wo,n,m)->float:
    """
    函数返回importance_sample_ggx的概率密度函数pdf
    pdf = distribution_ggx(nh_dot,m.a) * nh_dot
    """
    h = ti.math.normalize(wo + wi)
    nh_dot = ti.math.dot(n,h)
    return distribution_ggx(nh_dot,m.a) * nh_dot + 1e-6

@ti.func
def brdf(wi,wo,n,m,color):
    """
    BRDF : bidirectional reflective distribution function
    wi 入射光向量 wo出射光向量 n表面法线 m 材质数据
    """
    nl_dot = ti.math.dot(n,wi)
    kd = (1. - m.f) * (1.- m.metallic)
    h = ti.math.normalize(wo + wi)
    nv_dot = ti.math.dot(n,wo)
    nh_dot = ti.math.dot(n,h)
    d = distribution_ggx(nh_dot,m.a)
    g = geometry_smith(nv_dot,nl_dot,m.a)
    f = fresnel_schlick(nh_dot,m.f)
    specular = (d * g * f)/(4.* nv_dot *nl_dot + 1e-6)
    return kd * color/np.pi + specular

@ti.func
def lambert(wi,wo,n):
    return ti.math.dot(n,ti.math.normalize(wo))/np.pi

@ti.func
def lerp(v0,v1,t):
    """
    向量差值 t 0~1
    """
    return v0 + (v1-v0)*t
